# ------------------------------------------------------------------------------------------------
### Generate Figure 8: Distribution of synthetic control weights
# ------------------------------------------------------------------------------------------------

# Load data
load("../data/processed/LGBTQ_election_data.Rda")

# Filter data (exclude parliamentary election 2023 and Warsaw)
LGBTQ_election_data <- LGBTQ_election_data %>%
  filter(id_county != 1465)

# Define outcomes for analysis
outcomes <- list(
  turnout = "turnout",
  opposition_turnout = "non_pis_support", 
  government_turnout = "pis_support"
)

# Function to extract synthdid weights for a given outcome
extract_weights <- function(data, outcome_var) {
  sc_did <- data %>% select(id, eyear, !!sym(outcome_var), treat)
  setup <- panel.matrices(as.data.frame(as_tibble(sc_did)))
  estimates <- synthdid_estimate(setup$Y, setup$N0, setup$T0)
  weights <- synthdid_controls(estimates, mass = 1.001, weight.type = "omega")
  
  data.frame(
    id = rownames(weights),
    weight = weights[, "estimate 1"]
  )
}

# Extract weights for all outcomes
weights_list <- lapply(outcomes, function(outcome) {
  extract_weights(LGBTQ_election_data, outcome)
})

# Rename the weight columns in each data frame for clarity
weights_list$turnout <- weights_list$turnout %>% rename(w_turnout = weight)
weights_list$opposition_turnout <- weights_list$opposition_turnout %>% rename(w_opposition_turnout = weight)
weights_list$government_turnout <- weights_list$government_turnout %>% rename(w_government_turnout = weight)

# Load and prepare municipality boundaries
PL_gmina <- gisco_get_lau(country = "PL", year="2020") %>%
  mutate(id = paste0("", substr(LAU_ID, 5, 6), "", substr(LAU_ID, 10, nchar(LAU_ID)))) %>%
  left_join(LGBTQ_election_data %>% select(id, any_level_any_type) %>% unique(), by = "id") %>%
  left_join(weights_list$turnout, by = "id") %>%
  left_join(weights_list$opposition_turnout, by = "id") %>%
  left_join(weights_list$government_turnout, by = "id")

# Merge treated polygons for pattern efficiency
zones <- st_union(subset(PL_gmina, is.na(w_turnout)))

# Prepare data for plotting
PL_gmina_combined <- PL_gmina %>%
  gather("outcome", "weight", w_turnout:w_opposition_turnout:w_government_turnout) %>%
  mutate(
    outcome = factor(outcome, 
                     levels = c("w_turnout", "w_opposition_turnout", "w_government_turnout"), 
                     labels = c("Overall Turnout", "Opposition Turnout", "Government Turnout"))
  ) 

# Define fill scale for weights
fill_scale <- scale_fill_gradientn(
  colors = c("white", "darkblue"), 
  na.value = "transparent", 
  name = "sDiD weight",
  limits = c(0, round(max(PL_gmina_combined$weight, na.rm = TRUE), 3)),
  breaks = seq(0, round(max(PL_gmina_combined$weight, na.rm = TRUE), 3), 0.001),
  labels = seq(0, round(max(PL_gmina_combined$weight, na.rm = TRUE), 3), 0.001),
  guide = guide_colourbar(
    title.position = "top", 
    title.hjust = 0.5,
    text = element_text(size = 23),
    barwidth = 25
  )
)

# Generate plot
ggplot(PL_gmina_combined) +
  geom_sf(aes(fill = weight)) +
  fill_scale +
  geom_sf(data = zones, fill = "white") +
  facet_grid(outcome ~ ., switch = "y") +
  theme_bw() +
  theme(
    panel.grid = element_blank(),
    axis.text = element_blank(),
    axis.ticks = element_blank(),
    legend.position = "bottom",
    text = element_text(size = 23)
  )

# Save the plot
ggsave("../output/plots/figure8.pdf", width = 14, height = 18, dpi = 700)
